MeanMetricWrapper metric

@synalinks_export("synalinks.metrics.MeanMetricWrapper")
class MeanMetricWrapper(Mean):
    """Wrap a stateless metric function with the `Mean` metric.

    You could use this class to quickly build a mean metric from a function. The
    function needs to have the signature `fn(y_true, y_pred)` and return a
    per-sample reward array. `MeanMetricWrapper.result()` will return
    the average metric value across all samples seen so far.


    Example:

    ```python
    program.compile(
        metrics=[
            synalinks.metrics.MeanMetricWrapper(),
        ],
    )
    ```

    Args:
        fn (callable): The metric function to wrap, with signature
            `fn(y_true, y_pred, **kwargs)`.
        name (str): (Optional) string name of the metric instance.
        in_mask (list): (Optional) list of keys to keep to compute the metric.
        out_mask (list): (Optional) list of keys to remove to compute the metric.
        in_mask_pattern (str): (Optional) Regex pattern; fields whose names match
            are kept (combined with ``in_mask`` via OR).
        out_mask_pattern (str): (Optional) Regex pattern; fields whose names match
            are dropped (combined with ``out_mask`` via OR).
        **kwargs (keyword arguments): Keyword arguments to pass on to `fn`.
    """

    def __init__(
        self,
        fn,
        name=None,
        in_mask=None,
        out_mask=None,
        in_mask_pattern=None,
        out_mask_pattern=None,
        **kwargs,
    ):
        super().__init__(
            name=name,
            in_mask=in_mask,
            out_mask=out_mask,
            in_mask_pattern=in_mask_pattern,
            out_mask_pattern=out_mask_pattern,
        )
        self._fn = fn
        self._fn_kwargs = kwargs

        # If we are wrapping a Synalinks reward, register the metric's
        # direction as "up" (needs to be maximized during training).
        if (
            self._fn in rewards.ALL_OBJECTS
            or hasattr(self._fn, "__class__")
            and self._fn.__class__ in rewards.ALL_OBJECTS
        ):
            self.direction = "up"

    async def update_state(self, y_true, y_pred):
        y_pred = tree.map_structure(lambda x: ops.convert_to_json_data_model(x), y_pred)
        y_true = tree.map_structure(lambda x: ops.convert_to_json_data_model(x), y_true)
        if self.in_mask or self.in_mask_pattern:
            y_pred = tree.map_structure(
                lambda x: (
                    x.in_mask(mask=self.in_mask, pattern=self.in_mask_pattern)
                    if x is not None
                    else x
                ),
                y_pred,
            )
            y_true = tree.map_structure(
                lambda x: (
                    x.in_mask(mask=self.in_mask, pattern=self.in_mask_pattern)
                    if x is not None
                    else x
                ),
                y_true,
            )
        if self.out_mask or self.out_mask_pattern:
            y_pred = tree.map_structure(
                lambda x: (
                    x.out_mask(mask=self.out_mask, pattern=self.out_mask_pattern)
                    if x is not None
                    else x
                ),
                y_pred,
            )
            y_true = tree.map_structure(
                lambda x: (
                    x.out_mask(mask=self.out_mask, pattern=self.out_mask_pattern)
                    if x is not None
                    else x
                ),
                y_true,
            )
        values = await self._fn(y_true, y_pred, **self._fn_kwargs)
        return await super().update_state(values)

    def get_config(self):
        """Returns the serializable config of the metric."""
        base_config = super().get_config()
        config = {
            "fn": serialization_lib.serialize_synalinks_object(self._fn),
        }
        config.update(self._fn_kwargs)
        return {**base_config, **config}

    @classmethod
    def from_config(cls, config):
        fn = serialization_lib.deserialize_synalinks_object(config.pop("fn"))
        return cls(fn=fn, **config)